Verify each finding against the current code and only fix it if needed.

Inline comments:
In `@modelopt/torch/prune/plugins/mcore_minitron.py`:
- Around line 870-874: You’re overwriting TE modules’ return_layernorm_output
state unconditionally; instead capture each TELayerNormColumnParallelLinear’s
original return_layernorm_output before changing it (e.g., store in a dict keyed
by id(module) or attach a private attribute like _orig_return_layernorm_output
on the module) when you set it to True, and in the cleanup loop restore each
module’s original value rather than forcing False; apply the same
save-and-restore pattern for the other similar block referenced (lines
~999-1006) so original behavior is preserved after pruning/search.

---

Nitpick comments:
In `@modelopt/torch/prune/plugins/mcore_minitron.py`:
- Around line 918-935: The current set_local_activations_and_layer_scores method
assigns m._activations = local_activations[n] without validating the key, which
will raise a raw KeyError on checkpoint/model drift; update
set_local_activations_and_layer_scores to check if n is in local_activations
before assignment (iterate over self.model.named_modules()), and if missing
either raise a clear ValueError that includes the module name n and a summary of
available keys (e.g., list(local_activations.keys())) or log a descriptive
warning and skip restoring that module, so failures are diagnosable; reference
the method name set_local_activations_and_layer_scores and attributes
_activations, local_activations, and model.named_modules() when making the
change.

Verify each finding against the current code and only fix it if needed.

Inline comments:
In `@tests/gpu_megatron/torch/prune/plugins/test_mcore_mamba_minitron_pruning.py`:
- Around line 74-75: The tests set transformer_impl="transformer_engine" but
lack a guard, so add a Transformer Engine availability skip: either call
pytest.importorskip("megatron.core.extensions.transformer_engine") at module
scope in the test module (alongside the existing skip_if_no_mamba() call) or
implement a helper skip_if_no_transformer_engine() in
_test_utils/import_helper.py and invoke it next to skip_if_no_mamba(); ensure
the check runs before any model construction that uses transformer_impl to avoid
runtime failures.

---

Nitpick comments:
In `@tests/_test_utils/torch/megatron/models.py`:
- Around line 308-309: The Mamba helper currently treats any transformer_impl
value other than "transformer_engine" as "modelopt", allowing silent typos;
update the helper(s) that read the transformer_impl parameter (around the
transformer_impl = "modelopt" default and the code at the other occurrence) to
perform explicit validation: accept only the supported strings ("modelopt" and
"transformer_engine") and raise a ValueError with a clear message if an unknown
value is passed, rather than silently defaulting—adjust both occurrences (lines
near transformer_impl and the block at 390-393) to enforce this check.

Verify each finding against the current code and only fix it if needed.

Inline comments:
In `@modelopt/torch/prune/plugins/mcore_minitron.py`:
- Line 37: Wrap the module-level import of TELayerNormColumnParallelLinear in a
try/except that sets TELayerNormColumnParallelLinear = None and a HAS_TE boolean
flag (True on success, False on ImportError), then conditionally execute any
Transformer-Engine-specific logic only when HAS_TE is True; specifically guard
every place that references TELayerNormColumnParallelLinear (e.g., isinstance
checks and TE hook registrations) and the TE-specific hook registration blocks
so they run only if HAS_TE is True.

In `@modelopt/torch/utils/plugins/mbridge.py`:
- Around line 95-103: The else branch calls provider.num_moe_experts and
provider.qk_layernorm without defensive checks; update the branch around
transformer_layer_spec assignment (get_gpt_layer_with_transformer_engine_spec)
to either guard these attribute accesses with hasattr(provider,
"num_moe_experts") and hasattr(provider, "qk_layernorm") and handle missing
attributes (fallback values or raise a clear error), or add a concise
comment/docstring near the else branch documenting the assumption that provider
is a GPTModelProvider with those attributes (and cite the external megatron
source). Ensure references to provider.num_moe_experts, provider.qk_layernorm,
and get_gpt_layer_with_transformer_engine_spec are updated accordingly so future
callers see the validation or documented contract.

In `@tests/gpu_megatron/torch/nas/plugins/test_megatron_mamba_dynamic_modules.py`:
- Around line 75-76: The test hardcodes transformer_impl="transformer_engine"
but only calls skip_if_no_mamba(), so when Mamba exists and TE does not the test
will fail; add a new helper skip_if_no_transformer_engine() to
tests/_test_utils/import_helper.py (modeled on the TE guards used in
test_megatron.py around line ~964) and call skip_if_no_transformer_engine() at
the module level in
tests/gpu_megatron/torch/nas/plugins/test_megatron_mamba_dynamic_modules.py
alongside the existing skip_if_no_mamba() so the test is cleanly skipped when
Transformer Engine is unavailable.

---

Outside diff comments:
In `@modelopt/torch/prune/plugins/mcore_minitron.py`:
- Around line 477-513: The validation loop temporarily disables router flags and
mutates the model but only restores them on the happy path; make it
exception-safe by wrapping the candidate validation loop (the logic that
iterates top_k_candidates and calls _prune, eval_score, sample, and
saves/restores model.decoder.layers) in a try/finally so that in the finally
you: (1) re-enable every module in _routers_with_expert_bias by setting
m.enable_expert_bias = True, and (2) restore the model to the max subnet and
original layer numbering (use sample(self.model, sample_func=max) and reset
layer.layer_number from the saved start_layer_number and reassign
self.model.decoder.layers) to ensure cleanup runs even if _prune() or
eval_score() throws.

---

Duplicate comments:
In `@modelopt/torch/prune/plugins/mcore_minitron.py`:
- Around line 277-292: The cleanup flow currently forces every fused TE module's
return_layernorm_output to False and the registration path sets them to True,
which alters module behavior after pruning; update the code (in
registration/patching logic that flips fused TE modules and in cleanup()) to
record each module's original return_layernorm_output value (e.g., store a
per-module map keyed by module identity) and restore that exact original boolean
in cleanup() instead of hardcoding False, and ensure run_search() invokes
cleanup() inside a finally block so the original flags are restored even if an
exception occurs; refer to functions/methods cleanup(), run_search(), and the
return_layernorm_output flag on the fused TE modules when making the changes.

In `@tests/gpu_megatron/torch/prune/plugins/test_mcore_mamba_minitron_pruning.py`:
- Around line 74-75: The tests set transformer_impl="transformer_engine" but
only call skip_if_no_mamba(), which lets cases run when Mamba exists but
Transformer Engine does not; add a guard that imports or skips the TE module
before importing or instantiating the model/plugin (e.g., call
pytest.importorskip("megatron.core.extensions.transformer_engine") or create/use
a shared helper that does that) and apply the same change for the other
occurrences where transformer_impl="transformer_engine" is used; keep the
existing skip_if_no_mamba() but ensure the TE import-or-skip runs earlier to
prevent import failures.

Verify each finding against the current code and only fix it if needed.

Inline comments:
In `@modelopt/torch/nas/plugins/megatron.py`:
- Around line 390-403: The specialized TE parallel linear class registers only
"out_features" (via self._register_dynamic_attribute("out_features", ...)) but
not "in_features", causing stale dimension metadata after slicing; update the
class to also call self._register_dynamic_attribute("in_features", <callable>)
mirroring the parent _DynamicTEParallelLinear behavior (compute input dim from
mod.config.kv_channels and any active/group counts as appropriate), and ensure
"in_features" and "out_features" are kept in sync with the sliced weight/bias
(use the same pattern/closures used for "out_features" and for
bias/_get_weight/_get_bias/_get_ln_param to compute current dimensions).
- Around line 25-31: The module currently imports transformer_engine and the
Megatron TE extension classes (transformer_engine, TEColumnParallelLinear,
TEDotProductAttention, TELayerNormColumnParallelLinear, TERowParallelLinear)
unconditionally which breaks imports when TE is absent; wrap these imports in a
try/except and set a HAS_TE boolean (mirroring the HAS_MAMBA pattern already in
this file) so downstream registration of TE-specific features only occurs when
HAS_TE is True, and ensure any code that references those TE symbols is guarded
by the HAS_TE check similar to the pattern used in
modelopt/torch/quantization/plugins/megatron.py.

---

Outside diff comments:
In `@modelopt/torch/prune/plugins/mcore_minitron.py`:
- Around line 476-513: The loop temporarily disables module expert bias by
setting m.enable_expert_bias = False and collects them in
_routers_with_expert_bias, but if eval_score() or save_search_checkpoint()
throws the routers remain disabled; wrap the validation phase (the for candidate
in tqdm(...) loop and its inner pruning/eval/save logic) in a try/finally (or
surround the whole top-k validation block) and in the finally re-enable each
router by setting m.enable_expert_bias = True (using the existing
_routers_with_expert_bias list) so expert bias is always restored even on
exceptions.

---

Duplicate comments:
In `@modelopt/torch/prune/plugins/mcore_minitron.py`:
- Line 37: The module currently hard-imports TELayerNormColumnParallelLinear
which causes import failure when Transformer Engine (TE) is absent; wrap the
import in a try/except ImportError and set a module-level feature flag (e.g.,
_HAS_TE = True/False) accordingly, then guard all TE-specific logic/hook usage
(references to TELayerNormColumnParallelLinear and any TE-only hooks in the
plugin class in mcore_minitron.py) behind that flag so importing the module
succeeds even when TE is not installed.
- Around line 881-885: The cleanup loop currently forces
return_layernorm_output=False on every TELayerNormColumnParallelLinear in
self.model, mutating modules that may have originally been True; instead track
which specific modules you patched (e.g., store them in a list when you set
return_layernorm_output=True where you install hooks) and in the cleanup only
restore each tracked module's original value (store the original value
per-module when patching). Update the code that sets return_layernorm_output to
record (module, original_value) and replace the current iteration over
self.model.modules() with iteration over that recorded list to restore
original_value on the same TELayerNormColumnParallelLinear instances you
modified.

---

Nitpick comments:
In `@modelopt/torch/prune/plugins/mcore_minitron.py`:
- Around line 186-190: The checkpoint payload type hints are incorrect: adjust
the annotations for local_activations and layer_scores (and any related aliases
near sorted_layers and all_candidates_per_constraint) to match what is actually
stored — make local_activations a union type reflecting both nested dicts
(dict[str, dict[str, torch.Tensor]]) and flat dicts (dict[str, torch.Tensor])
produced by the hidden-size collector and SequentialMLP collector, and change
layer_scores from dict[int, torch.Tensor] to dict[int, float] (since .item() is
stored); ensure sorted_layers remains Optional[List[int]] (1-indexed) and
all_candidates_per_constraint stays dict[float, list[CandidateSubnet]]; update
any type aliases or function signatures that construct or consume these payloads
so mypy and checkpoint serializers see the real shapes (search for
local_activations, layer_scores, sorted_layers, all_candidates_per_constraint,
CandidateSubnet to locate all usages).